# -*- coding: utf-8 -*-
"""
Created on Sat Jun  9 12:43:20 2012

@author: nzhao
"""

from traits.api import HasTraits, Enum, Float, Str, Instance

import numpy as np

import phy.general.Constant as co
from phy.spin.AbstractSpinSystem import AbstractSpinSystem
from math_qm.DensityMatrix import DensityMatrix_1Qubit
from math_qm.SpinHamiltonian import SingleSpin_Hamiltonian
from math_qm.vector import vector

dft_b = vector()

class SingleSpin(AbstractSpinSystem, HasTraits):
    spin_element = Enum('NV', 'E', 'C13', 'N14', 'N15', 'H')
    spin_type    = Enum('Central', 'Bath')
    spin_mark    = Str('None')
    spin_relative_Gamma = Float
    
    x = Float(0.0)
    y = Float(0.0)
    z = Float(0.0)

    state_I = Float(1.0)
    state_X = Float(0.0)
    state_Y = Float(0.0)
    state_Z = Float(0.0)
    
    dist  = Float
    
    def __init__(self, element = 'NV', spin_type = 'Central', spin_status = 'On', state = 'thermal', x = 0.0, y = 0.0, z = 0.0):
        self.spin_element = element
        self.spin_relative_Gamma = co.SpinDictionary[self.spin_element]['RelativeGyromagneticRatio']

        self.spin_type = spin_type
        self.spin_status = spin_status

        self.x, self.y, self.z = x, y, z
        self.dist = (x**2 + y**2 + z**2)**0.5
        
        self.state = DensityMatrix_1Qubit()
        self.state.on_trait_change(self.update_state_vector, 'state_I, state_X, state_Y, state_Z')
        
        self.hamiltonian = SingleSpin_Hamiltonian(spin_element = element)
        
    def set_hamiltonian(self, magnetic_field = None):
        if not isinstance(magnetic_field, vector):
            print 'Unknown magnetic field. Use defalt(zero) magnetic field.'
            magnetic_field = dft_b
                    
        self.hamiltonian.calculate_hamiltonian(magnetic_field)
        
        return self.hamiltonian.matrix
    
    def expectation_value(self, op_lst, state=0):
        wavefunc = self.calculate_eigs() [1]
        res = np.array([np.dot(np.conjugate(wavefunc[:, state]), np.dot(op, wavefunc[:, state])) for op in op_lst])
        return res
    
    def postion(self):
        return np.array([self.x, self.y, self.z])
    
    def update_state_vector(self):
        self.state_I, self.state_X, self.state_Y, self.state_Z = self.state.get_vector()
    def _state_I_changed(self):
        self.state.set_vector(self.state_I, self.state_X, self.state_Y, self.state_Z)
    def _state_X_changed(self):
        self.state.set_vector(self.state_I, self.state_X, self.state_Y, self.state_Z)    
    def _state_Y_changed(self):
        self.state.set_vector(self.state_I, self.state_X, self.state_Y, self.state_Z)    
    def _state_Z_changed(self):
        self.state.set_vector(self.state_I, self.state_X, self.state_Y, self.state_Z)

    def _spin_element_changed(self):
        self.spin_relative_Gamma = co.SpinDictionary[self.spin_element]['RelativeGyromagneticRatio']


if __name__ == '__main__':
    spin = SingleSpin()
    print 'before change state_I=', spin.state_I, spin.state_X, spin.state_Y, spin.state_Z
    print 'before change state_I=', spin.state.state_I, spin.state.state_X, spin.state.state_Y, spin.state.state_Z
    spin.state_I = 0.5
    print 'after change state_I=', spin.state_I, spin.state_X, spin.state_Y, spin.state_Z
    print 'after change state_I=', spin.state.state_I, spin.state.state_X, spin.state.state_Y, spin.state.state_Z
    
    print 
    
    spin.state.state_Z = 0.1
    print 'after change state.state_Z=', spin.state_I, spin.state_X, spin.state_Y, spin.state_Z
    print 'after change satte.state_Z=', spin.state.state_I, spin.state.state_X, spin.state.state_Y, spin.state.state_Z